#include "UGV.h"

namespace iusc
{
    void UGV::initialize()
    {
        x = m_param.x;
        y = m_param.y;
        z = m_param.z;
        m_param.vx = 0.0;
        m_param.vy = 0.0;
        m_param.ax = 0.0;
        m_param.ay = 0.0;
        m_id = CustomObjectID::UGV_VEHICLE;
    }

    /**
     * @brief 当前仅支持直行
     * 
     * @param dt 
     */
    void UGV::update(double dt)
    {
        m_param.x += m_param.vx*dt + 0.5*m_param.ax*dt*dt;
        m_param.y += m_param.vy*dt + 0.5*m_param.ay*dt*dt;
        m_param.vx += m_param.ax*dt;
        m_param.vy += m_param.ay*dt;
        z = m_param.z;
        x = m_param.x;
        y = m_param.y;
    }

    void UGV::getPose(geometry_msgs::PoseStamped& pose) const
    {
        pose.header.frame_id = m_param.link;
        pose.pose.position.x = x;
        pose.pose.position.y = y;
        pose.pose.position.z = z;
        pose.pose.orientation.w = 1.0;
        pose.pose.orientation.y = 0;
        pose.pose.orientation.x = 0;
        pose.pose.orientation.z = 0;
    }

    // === UGV Formation ====

    void UGVFormation::initialize()
    {
        m_id = CustomObjectID::UGV_FORMATION;
        // 加载参数文件
        for (auto u:m_ugvObjects)
        {
            u->initialize();
        }
    }

    void UGVFormation::update(double dt)
    {
        if (m_stopFlag || !m_enableFlag){  // 如果停止了啥都不做
            return ;
        }
        if (m_stage1Flag){  // 先进行编队集结
            updateStage1(dt);
        }
        else {  // 已经编好队了,直接根据规律运动
            updateStage2(dt);
        }       
        for (auto u:m_ugvObjects)
        {
            u->update(dt);
        }
    }

    /**
     * @brief 对于编队来说，这里会返回编队的中心位置
     * 
     * @param pose 
     */
    void UGVFormation::getPose(geometry_msgs::PoseStamped& pose) const
    {
        pose.header.frame_id = "world";
        pose.pose.position.x = m_formationParams.x;
        pose.pose.position.y = m_formationParams.y;
        pose.pose.position.z = m_formationParams.z;
        pose.pose.orientation.w = 1.0;
        pose.pose.orientation.y = 0;
        pose.pose.orientation.x = 0;
        pose.pose.orientation.z = 0;
    }

    /**
     * @brief 添加无人车，仅在初始化的时候用
     * 
     * @param ugv 
     * @return true 
     * @return false 
     */
    bool UGVFormation::addUGV(UGV::Ptr ugv)
    {
        bool existFlag = false;
        for (auto u:m_ugvObjects)
        {
            if (ugv->ID() == u->ID()){
                existFlag = true;
                break;
            }
        }
        if (!existFlag){
            m_ugvObjects.push_back(ugv);
        }
        return !existFlag;  // 返回是否添加成功
    }

    /**
     * @brief 获取所有无人车的位置
     * 
     * @param ugvPoses 
     */
    void UGVFormation::getGroundTruth(std::vector<geometry_msgs::PoseStamped>& ugvPoses) const
    {
        std::vector<geometry_msgs::PoseStamped> res;
        for (auto u:m_ugvObjects)
        {
            geometry_msgs::PoseStamped pose;
            u->getPose(pose);
            res.emplace_back(pose);
        }
        res.swap(ugvPoses);
    }

    /**
     * @brief 将所有无人车放到初始位置，编队重整，重新运行
     * 
     * @param param 
     */
    void UGVFormation::reset(UGVFormationParams& param)
    {
        m_enableFlag = false;
        m_stopFlag = false;
        m_stage1Flag = true;

        m_formationRelPoses.clear();  // 编队需要重新确定
        std::map<int,UGV::Ptr> newFormationMap;
        m_formationMap.swap(newFormationMap);
        m_trajT.clear();
        m_trajVel.clear();

        // m_formationParams.x = m_formationInitX;
        // m_formationParams.y = m_formationInitY;
        // m_formationParams.vx = 0;
        // m_formationParams.vy = 0;
        setParam(m_formationParams);

        for (int i=0;i<m_ugvObjects.size();i++){
            auto ugv = m_ugvObjects[i];
            UGVParams p;
            ugv->getParam(p);
            p.x = m_initializeUGVPoses[i].x;
            p.y = m_initializeUGVPoses[i].y;
            p.z = m_initializeUGVPoses[i].z;
            p.vx = 0;
            p.vy = 0;
            p.ax = 0;
            p.ay = 0;
            ugv->setParam(p);
            ugv->update(0);
        }
        m_trajExecTime = 0.0;
    }

    void UGVFormation::start(){
        m_enableFlag = true;  // 允许开始运动
    }

    void UGVFormation::stop()
    {  // 停止运动
        m_stopFlag = true;
    }

    /**
     * @brief 编队的第一阶段，负责分配任务，然后集结运动
     * 
     */
    void UGVFormation::updateStage1(double dt)
    {
        if (m_formationMap.empty()){
            assignUGVs();  // 先获取编队相对位置，然后变换到世界下求解
        }
        else {  // 运动
            moveStage1UGVs(dt);
        }
        bool ret = judgeStage1Finished();
        // 如果结束则需要生成轨迹
        if (ret){
            m_stage1Flag = false;
            std::cout<<"=== Generate Formation Trajectory ==="<<std::endl;
            if (m_formationParams.moveType == 0){
                double R = 4*m_formationParams.distance;
                double omega = 1.0;
                formation::generateCircleTraj(R,omega,m_trajT,m_trajVel);
            }
            else if(m_formationParams.moveType == 1){
                double L = 4*m_formationParams.distance;
                double v = m_formationParams.vmax;
                formation::generateRectTraj(L,v,m_trajT,m_trajVel);
            }
            else {
                double v = m_formationParams.vmax;
                formation::generateRandomTraj(v,m_trajT,m_trajVel);
            }
            m_totalTrajTime = 0.0;
            for (int i=0;i<m_trajT.size();i++){
                m_totalTrajTime += m_trajT[i];
                std::cout<<"Piece "<<i<<": T:"<<m_trajT[i]<<", V:"<<m_trajVel[i].x<<","<<m_trajVel[i].y<<std::endl;
            }
        }
    }

    /**
     * @brief 获取编队相对位置，然后直接设置无人车位置,后面再改
     * 
     */
    void UGVFormation::assignUGVs()
    {
        double distance = m_formationParams.distance;
        formation::FormationType type;
        if (m_formationParams.formationType == 0){
            type = formation::TRIANGLE;
        }
        else if (m_formationParams.formationType == 1){
            type = formation::RECTANGLE;
        }
        else {
            type = formation::CIRCLE;
        }
        formation::getFormationRelPose(m_formationRelPoses,
            m_formationParams.distance,type);
        std::cout<<"Get Formation Rel Poses ,Size:"<<m_formationRelPoses.size()<<std::endl;
        std::vector<formation::Point> formationAbsPoses;  // 在世界坐标系下的位置
        std::vector<formation::Point> ugvInitPoses;  // 无人车初始位姿
        const size_t N = m_formationRelPoses.size();
        for (int i=0;i<m_formationRelPoses.size();i++){
            formation::Point p = m_formationRelPoses[i];
            p.x += m_formationInitX;
            p.y += m_formationInitY;
            p.z += m_formationParams.z;
            formationAbsPoses.emplace_back(p);
        }

        std::vector<UGV::Ptr> landedUGVs;
        for (auto ugv:m_ugvObjects){
            if (ugv->Landed()){
                landedUGVs.push_back(ugv);
                UGVParams p;
                ugv->getParam(p);
                formation::Point pt;
                pt.x = p.x;
                pt.y = p.y;
                pt.z = p.z;
                ugvInitPoses.emplace_back(pt);
            }
            else {
                m_randomUGVs.push_back(ugv);
            }
        }

        std::cout<<"Landed UGV Size:"<<ugvInitPoses.size()<<std::endl;
        // 构建代价矩阵求解指派问题
        std::vector<std::vector<double>> costs;
        double maxDis = 0.0;
        for (int i=0;i<N;i++){
            std::vector<double> iDist;
            iDist.resize(N);
            auto pi = ugvInitPoses[i];
            for (int j=0;j<N;j++){
                auto pj = formationAbsPoses[j];
                double dis = (pi.x - pj.x)*(pi.x - pj.x) + (pi.y-pj.y)*(pi.y - pj.y);
                if (dis > maxDis){maxDis = dis;}
                iDist[j] = -dis;
            }
            costs.emplace_back(iDist);
        }
        std::cout<<"Build Assignmet Problem!"<<std::endl;
        std::cout<<"Cost Matrix :"<<std::endl;
        for (int i=0;i<N;i++){
            for (int j=0;j<N;j++){
                costs[i][j] /= maxDis;  // 归一化一下
                std::cout<<costs[i][j]<<"\t";
            }
            std::cout<<std::endl;
        }
        
        std::vector<int> row4col,col4row;
        double gain;
        bool ret = assign2dAuction(costs,col4row,row4col,gain,100);
        std::cout<< "Solve Finished,Assign Result : "<<ret<<",cost:"<<gain<<std::endl;
        for (int i=0;i<landedUGVs.size();i++){
            UGVParams p;
            landedUGVs[i]->getParam(p);
            int fInd = col4row[i];
            if (fInd == -1){
                std::cout<<"UGV i:"<<landedUGVs[i]->ID()<<",NO Assigned!!!"<<std::endl;
                continue;}
            p.x = formationAbsPoses[fInd].x;
            p.y = formationAbsPoses[fInd].y;
            landedUGVs[i]->setParam(p);  // 
            std::cout<<"UGV i:"<<landedUGVs[i]->ID()<<",formation Ind :"<<fInd<<"Req Pos:"<<p.x<<","<<p.y<<std::endl;
            m_formationMap[fInd] = landedUGVs[i];
        }
    }

    void UGVFormation::moveStage1UGVs(double dt){
        return ;
    }

    bool UGVFormation::judgeStage1Finished() {
        return true;
    }

    /**
     * @brief 编队的第二阶段
     * @note 直接写一个简单版，所有车都跟编队的速度走
     */
    void UGVFormation::updateStage2(double dt)
    {
        m_trajExecTime += dt;
        if (m_trajExecTime > m_totalTrajTime){
            m_trajExecTime -= m_totalTrajTime;
        }
        double t = m_trajExecTime;
        for (int i=0;i<m_trajT.size();i++){
            if (t > m_trajT[i]){
                t -= m_trajT[i];
            }
            else {
                m_formationParams.vx = m_trajVel[i].x;
                m_formationParams.vy = m_trajVel[i].y;
                break;
            }
        }
        for (auto it = m_formationMap.begin();it != m_formationMap.end();it++){
            UGVParams p;
            it->second->getParam(p);
            p.vx = m_formationParams.vx;
            p.vy = m_formationParams.vy;
            it->second->setParam(p);
        }
        // 其他车随机生成
        // std::random_device rd;
        // std::mt19937 gen(rd());
        // double noise_v = std::min(m_formationParams.vx,m_formationParams.vy);
        // std::uniform_real_distribution<> dis(-noise_v,noise_v);
        // double thre = m_formationParams.distance;
        // for (auto& ugv:m_randomUGVs){
        //     double x,y;
        //     double vx,vy;
        //     geometry_msgs::PoseStamped p;
        //     ugv->getPose(p);
        //     UGVParams par;
        //     ugv->getParam(par);
        //     x = p.pose.position.x;
        //     y = p.pose.position.y;
        //     vx = par.vx;
        //     vy = par.vy;  // 先用之前的速度
        //     int iter = 0;
        //     double max_dis,max_vx,max_vy;
        //     max_dis = 0.0;  // 如果还是有碰撞选一个集群距离最大的
        //     max_vx = vx;  // 最大距离对应的速度(针对超过迭代数的情况)
        //     max_vy = vy;
        //     double form_dis;  // 距离集群的最小距离
        //     while (checkCollision(x,y,vx,vy,dt,thre,form_dis) && iter < 10){
        //         if (form_dis > max_dis){
        //             max_dis = form_dis;
        //             max_vx = vx;
        //             max_vy = vy;
        //         }
        //         // vx = m_formationParams.vx + dis(gen);
        //         // vy = m_formationParams.vy + dis(gen);
        //         vx = dis(gen);
        //         vy = dis(gen);
        //         iter++;
        //     }
            
        //     par.vx = max_vx;
        //     par.vy = max_vy;
        //     ugv->setParam(par);
        // }
    }

    /**
     * @brief 检测速度是否会和编队发生碰撞
     * 
     * @param x 当前车x
     * @param y 当前车y
     * @param vx 期望速度
     * @param vy 
     * @param dt 预测时间
     * @param threshold 最近距离
     * @param dis 距离集群的最小距离
     * @return true 
     * @return false 
     */
    bool UGVFormation::checkCollision(double x,double y,double vx,double vy,double dt,double threshold,double& dis)
    {
        bool ret = false;
        double nx,ny;
        nx = x + vx*dt;
        ny = y + vy*dt;
        dis = 1e6;
        for (int i=0;i<m_ugvObjects.size();i++){
            const auto& ugv = m_ugvObjects[i];
            geometry_msgs::PoseStamped p;
            ugv->getPose(p);
            if (fabs(x-p.pose.position.x) < 1e-2 && fabs(y-p.pose.position.y) < 1e-2){
                continue;
            }
            double distance = fabs(nx-p.pose.position.x) < threshold && fabs(ny-p.pose.position.y);
            if (distance < dis){
                dis = distance;
            }
            if (distance < threshold)
            {
                ret = true;
            }
        }
        return ret;
    } 

}